# Copyright 2022 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""MulQuant."""
from __future__ import absolute_import

from mindspore.ops import operations as P
from mindspore.common.dtype import QuantDtype
from mindspore.nn.cell import Cell
from .fake_quant_with_min_max_observer import quant_config_default


class MulQuant(Cell):
    r"""
    Adds fake quantized operation after `Mul` operation.

    This part is a more detailed overview of `Mul` operation. For more details about Quantization,
    please refer to the implementation of class of `FakeQuantWithMinMaxObserver`,
    :class:`mindspore.nn.FakeQuantWithMinMaxObserver`.

    Args:
        ema_decay (float): Exponential Moving Average algorithm parameter. Default: 0.999.
        quant_config (QuantConfig): Configures the types of quant observer and quant settings of weight and
            activation. Note that, QuantConfig is a special namedtuple, which is designed for quantization
            and can be generated by :func:`mindspore.compression.quant.create_quant_config` method.
            Default: QuantConfig with both items set to default :class:`FakeQuantWithMinMaxObserver`.
        quant_dtype (QuantDtype): Specifies the FakeQuant datatype. Default: QuantDtype.INT8.

    Inputs:
        - **x1** (Tensor) - The first tensor of MulQuant. The input dimension is preferably 2D or 4D.
        - **x2** (Tensor) - The second tensor of MulQuant. Has the same shape with `x1`.

    Outputs:
        Tensor, with the same type and shape as the `x1`.

    Raises:
        TypeError: If `ema_decay` is not a float.
        ValueError: If the shape of `x2` is different with `x1`.

    Supported Platforms:
        ``Ascend`` ``GPU``

    Examples:
        >>> import numpy as np
        >>> import mindspore
        >>> from mindspore.compression import quant
        >>> from mindspore import Tensor, nn
        >>> qconfig = quant.create_quant_config()
        >>> mul_quant = nn.MulQuant(quant_config=qconfig)
        >>> x1 = Tensor(np.array([[1, 2, 1], [-2, 0, -1]]), mindspore.float32)
        >>> x2 = Tensor(np.ones((2, 3)) * 2, mindspore.float32)
        >>> output = mul_quant(x1, x2)
        >>> print(output)
        [[ 1.9764705  4.0000005  1.9764705]
         [-4.         0.        -1.9764705]]
    """

    def __init__(self,
                 ema_decay=0.999,
                 quant_config=quant_config_default,
                 quant_dtype=QuantDtype.INT8):
        """Initialize MulQuant."""
        super(MulQuant, self).__init__()
        self.fake_quant_act = quant_config.activation(min_init=-6,
                                                      max_init=6,
                                                      ema=True,
                                                      ema_decay=ema_decay,
                                                      quant_dtype=quant_dtype)
        self.mul = P.Mul()

    def construct(self, x1, x2):
        """construct."""
        x = self.mul(x1, x2)
        x = self.fake_quant_act(x)
        return x
